﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using Ast.Core.Services.Extensions;
using Ast.Framework;
using Ast.Framework.Caching;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface ISsdtService : IDependency
    {
        IntPtr Ssdt { get; }
        IntPtr SsdtShadow { get; }

        bool FindAddress();
        void ListFunctions(IntPtr ssdt, Func<SsdtService.SsdtFunction, bool> ssdtProcFunc);
    }

    public class SsdtService : ISsdtService
    {
        private readonly IPspCidTableService _pspCidTableService;
        private readonly IPsLoadedModuleListService _psLoadedModuleListService;
        private readonly IDeviceController _deviceController;
        private readonly ICacheManager _cacheManager;
        private readonly ISymbolService _symbolService;

        public SsdtService(IPspCidTableService pspCidTableService, IPsLoadedModuleListService psLoadedModuleListService, IDeviceController deviceController, ICacheManager cacheManager, ISymbolService symbolService)
        {
            _pspCidTableService = pspCidTableService;
            _psLoadedModuleListService = psLoadedModuleListService;
            _deviceController = deviceController;
            _cacheManager = cacheManager;
            _symbolService = symbolService;

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public IntPtr Ssdt { get; private set; }
        public IntPtr SsdtShadow { get; private set; }

        public bool FindAddress()
        {
            Ssdt = SsdtShadow = IntPtr.Zero; // reset

            _pspCidTableService.ListPspCidTableObjects((objPtr, objType) =>
            {
                if (objType == _pspCidTableService.ProcessType) return true; // continue search if it is a process

                // search by interating all threads
                var thread = _deviceController.ReadAs<EThread>(objPtr);
                if (thread.Win32Thread == 0)
                {
                    if (Ssdt == IntPtr.Zero)
                    {
                        Ssdt = thread.ServiceTable;
                        Logger.InfoFormat("found ssdt 0x{0:x}", Ssdt.ToInt32());
                    }
                }
                else
                {
                    if (SsdtShadow == IntPtr.Zero)
                    {
                        SsdtShadow = thread.ServiceTable;
                        Logger.InfoFormat("found ssdt shadow 0x{0:x}", SsdtShadow.ToInt32());
                    }
                }

                return Ssdt == IntPtr.Zero || SsdtShadow == IntPtr.Zero; // continue search(returns true) if one of them not found yet
            });

            Logger.Info("ssdt and ssdt shadow search done");
            return Ssdt != IntPtr.Zero && SsdtShadow != IntPtr.Zero;
        }

        public void ListFunctions(IntPtr ssdt, Func<SsdtFunction, bool> ssdtProcFunc)
        {
            var tblDesc = _deviceController.ReadAs<KServiceTableDescriptor>(ssdt);
            var serviceTables = new List<KSystemServiceTable>(2);
            if (tblDesc.Ntoskrnl.ServiceTableBase != IntPtr.Zero) serviceTables.Add(tblDesc.Ntoskrnl);
            if (tblDesc.Win32k.ServiceTableBase != IntPtr.Zero) serviceTables.Add(tblDesc.Win32k);

            var kernelModules = _cacheManager.Get("kernelModules", context => _psLoadedModuleListService.GetAllModules());

            var goOn = true;
            foreach (var serviceTable in serviceTables)
            {
                for (var i = 0; i < serviceTable.NumberOfService; i++)
                {
                    var funcAdr = (uint) _deviceController.ReadAsIntPtr(serviceTable.ServiceTableBase + i*4).Value;
                    var kernelModule = kernelModules.FirstOrDefault(module => funcAdr >= module.StartAddress && funcAdr <= module.EndAddress);
                    if (kernelModule == null)
                    {
                        Logger.InfoFormat("not found module for function 0x{0:x}", funcAdr);
                        continue;
                    }

                    // symbol
                    var symbolBaseAddress = _symbolService.LoadModule(kernelModule.Name, null, kernelModule.StartAddress);
                    if (symbolBaseAddress == 0)
                    {
                        Logger.ErrorFormat("failed to load module symbol {0}", kernelModule.Name);
                        break;
                    }

                    var offset = funcAdr - kernelModule.StartAddress;
                    var symbolAdr = offset + symbolBaseAddress;
                    var symbolInfo = _symbolService.LoadSymbol(symbolAdr);
                    if (symbolInfo.HasValue)
                    {
                        if (ssdtProcFunc != null)
                        {
                            goOn = ssdtProcFunc(new SsdtFunction
                            {
                                Module = kernelModule,
                                ServiceTableBase = (uint) serviceTable.ServiceTableBase.ToInt32(),
                                FunctionIndex = i,
                                FunctionAddress = funcAdr,
                                FunctionName = symbolInfo.Value.Name
                            });
                            if (!goOn) break;
                        }
                    }
                    else
                    {
                        Logger.InfoFormat("not found symbol for function 0x{0:x}", funcAdr);
                    }
                }

                if (!goOn) break;
            }
        }

        public class SsdtFunction
        {
            public KernelModule Module { get; set; }
            public uint ServiceTableBase { get; set; }
            public int FunctionIndex { get; set; }
            public string FunctionName { get; set; }
            public uint FunctionAddress { get; set; }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct KSystemServiceTable
        {
            public IntPtr ServiceTableBase;
            public IntPtr ServiceCounterTableBase;
            public uint NumberOfService;
            public IntPtr ParamTableBase;
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto)]
        public struct KServiceTableDescriptor
        {
            [MarshalAs(UnmanagedType.Struct)]
            public KSystemServiceTable Ntoskrnl;

            [MarshalAs(UnmanagedType.Struct)]
            public KSystemServiceTable Win32k;
        }
    }
}